#include "QNBLockCtrlAltDel.h"

#include <tlhelp32.h>

extern "C"
{

#include <Windows.h>
#include <WinBase.h>


}

typedef LONG (NTAPI * NtSuspendProcess)(IN HANDLE);
typedef LONG (NTAPI * NtResumeProcess)(IN HANDLE);

QNBLockCtrlAltDel::QNBLockCtrlAltDel()
{
}

DWORD QNBLockCtrlAltDel::GetProcessIdFromName(QString Name, LPTSTR szID)
{
    PROCESSENTRY32 pe;
    DWORD   id   =   0;
    HANDLE   hsnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS,0);
    pe.dwSize   =   sizeof(PROCESSENTRY32);


    if(!Process32First(hsnapshot,&pe))
        return   0;
    do
    {
        pe.dwSize   =   sizeof(PROCESSENTRY32);
        if(Process32Next(hsnapshot,&pe)==false)
        break;
        QString str = QString::fromWCharArray( pe.szExeFile );
        if(Name.compare( str ) == 0)
        {
            id  =  pe.th32ProcessID;
            break;
        }


    } while(1);
    CloseHandle(hsnapshot);

    if(szID == NULL) {
        return id;
    }

    char lpID[10];
    itoa(id, lpID, 10);
    strcpy((char*)szID, lpID);
    return   id;
}

bool QNBLockCtrlAltDel::EnableDebugPrivilege(bool bEnableDebugPrivilege)
{
#if 1
    HANDLE hToken;
    TOKEN_PRIVILEGES tp;
    LUID luid;


    if(!::OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &hToken))
    {
        //::MessageBox(this->GetSafeHwnd(), GET_TOKEN_ERROR, MSG_BOX_TITLE, MB_OK);
        return FALSE;
    }


    if(!::LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &luid))
    {
        // ::MessageBox(this->GetSafeHwnd(), GET_PRIVILEGE_VALUE_ERROR, MSG_BOX_TITLE, MB_OK);
        CloseHandle(hToken);
        return FALSE;
    }

    tp.PrivilegeCount = 1;
    tp.Privileges[0].Luid = luid;
    if(bEnableDebugPrivilege)
    {
        tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
    }
    else
    {
        tp.Privileges[0].Attributes = 0;
    }


    if(!::AdjustTokenPrivileges(hToken, FALSE, &tp, sizeof(tp), NULL, NULL))
    {
    // ::MessageBox(this->GetSafeHwnd(), ADJUST_PRIVILEGE_ERROR, MSG_BOX_TITLE, MB_OK);
    CloseHandle(hToken);
    return FALSE;
    }


    CloseHandle(hToken);
    if(GetLastError() == ERROR_NOT_ALL_ASSIGNED)
    {
    // ::MessageBox(this->GetSafeHwnd(), ENABLE_DEBUG_ERROR, MSG_BOX_TITLE, MB_OK);
    return FALSE;
    }

    #endif
    return TRUE;
}


//winLogon.exe进程挂起与恢复
bool QNBLockCtrlAltDel::TaskScreen(bool type)
{
    bool bRet = EnableDebugPrivilege(true);//提权
    if( !bRet )
    {
        return false ;
    }

    LPTSTR csPath = NULL;
    QString pName = "winlogon.exe";
    DWORD dwId;
    dwId = GetProcessIdFromName(pName,csPath);//获取进程ID
    HANDLE hProcess;
    if (0 != dwId)
    {
        hProcess = OpenProcess(PROCESS_ALL_ACCESS,FALSE,dwId);
        if(NULL != hProcess)
        {

            //动态加载dll库
            HMODULE h_module=LoadLibrary(L"ntdll.dll");
            if( !type )
            {
                //进程挂起
                NtSuspendProcess m_NtSuspendProcess = (NtSuspendProcess)GetProcAddress(h_module,"NtSuspendProcess");
                DWORD dwRet = m_NtSuspendProcess(hProcess);
                DWORD dwErr = GetLastError();
            }
            else
            {    //进程恢复
                NtResumeProcess m_NtResumeProcess = (NtResumeProcess)GetProcAddress(h_module,"NtResumeProcess");
                DWORD dwRet = m_NtResumeProcess(hProcess);
                DWORD dwErr = GetLastError();
            }

            //采放动态库
            FreeLibrary(h_module);
        }
        else
        {
            return false;
        }
    }
    else
    {
        return false;
    }

    return true;
}
